use std::{
    collections::HashMap,
    io::Write,
    net::{TcpListener, TcpStream},
    sync::LazyLock,
};

use clap::Parser;
use clickhouse::Client;
use prometheus::{Encoder, Registry, TextEncoder};

static DEFAULT_REGISTRY: LazyLock<Registry> =
    LazyLock::new(|| Registry::new_custom(Some("clickhouse".to_string()), None).unwrap());

#[tokio::main]
async fn main() {
    let config = Config::parse();

    {
        tokio::spawn(async move {
            let mut exporter = Exporter::new(config);
            exporter.register().await.unwrap();

            if let Err(e) = exporter.collect().await {
                println!("{e}");
            }
        });
    }

    let listener = TcpListener::bind("0.0.0.0:9300").unwrap();
    println!("Server started at http://0.0.0.0:9300");

    for stream in listener.incoming() {
        if let Ok(stream) = stream {
            if let Err(e) = handle_request(stream) {
                println!("{e}");
            }
        }
    }
}

#[derive(clap::Parser, Debug, Clone)]
struct Config {
    /// Clickhouse host
    #[arg(long, default_value = "0.0.0.0")]
    host: String,

    /// Clickhouse port
    #[arg(long, default_value = "8123")]
    port: i32,

    /// Clickhouse username
    #[arg(long, default_value = "default")]
    username: String,

    /// Clickhouse password
    #[arg(long, default_value = "default")]
    password: String,
}

fn handle_request(mut stream: TcpStream) -> anyhow::Result<()> {
    let metric_family = DEFAULT_REGISTRY.gather();

    let encoder = TextEncoder::new();
    let mut buffer = Vec::new();

    encoder.encode(&metric_family, &mut buffer)?;

    let header = format!(
        "HTTP/1.1 200 OK\r\nContent-Type: {}\r\nContent-Length: {}\r\n\r\n",
        encoder.format_type(),
        buffer.len()
    );

    stream.write(header.as_bytes())?;
    stream.write(&buffer)?;
    stream.flush()?;

    Ok(())
}

struct Exporter {
    client: Client,
    metric_map: HashMap<String, prometheus::Gauge>,
    system_metric_sql: String,
    system_events_sql: String,
}

impl Exporter {
    fn new(config: Config) -> Exporter {
        let url = format!("http://{}:{}", config.host, config.port);

        let client = Client::default()
            .with_url(url)
            .with_user(config.username)
            .with_password(config.password);

        Exporter {
            client,
            metric_map: HashMap::new(),
            system_metric_sql: "select metric,value,description from system.metrics".to_string(),
            system_events_sql: "select event,value,description from system.events".to_string(),
        }
    }

    async fn register(&mut self) -> anyhow::Result<()> {
        {
            let rows = self
                .client
                .query(&self.system_events_sql)
                .fetch_all::<(String, u64, String)>()
                .await?;

            for row in rows {
                let Ok(metric) = prometheus::Gauge::new(&row.0, &row.2) else {
                    println!("{:?}", row);
                    continue;
                };

                DEFAULT_REGISTRY.register(Box::new(metric.clone())).ok();
                self.metric_map.insert(row.0, metric);
            }
        }

        {
            let rows = self
                .client
                .query(&self.system_metric_sql)
                .fetch_all::<(String, i64, String)>()
                .await?;

            for row in rows {
                let Ok(metric) = prometheus::Gauge::new(&row.0, &row.2) else {
                    println!("{:?}", row);
                    continue;
                };

                DEFAULT_REGISTRY.register(Box::new(metric.clone())).ok();
                self.metric_map.insert(row.0, metric);
            }
        }

        Ok(())
    }

    async fn collect(&mut self) -> anyhow::Result<()> {
        loop {
            self.collect_metric().await?;
            self.collect_events().await?;

            tokio::time::sleep(std::time::Duration::from_secs(15)).await;
        }
    }

    async fn collect_events(&mut self) -> anyhow::Result<()> {
        let rows = self
            .client
            .query(&self.system_events_sql)
            .fetch_all::<(String, u64, String)>()
            .await?;

        for row in rows {
            if let Some(m) = self.metric_map.get_mut(&row.0) {
                m.set(row.1 as f64);
                continue;
            }

            let metric = prometheus::Gauge::new(&row.0, row.2)?;
            metric.set(row.1 as f64);

            DEFAULT_REGISTRY.register(Box::new(metric.clone())).ok();
            self.metric_map.insert(row.0, metric);
        }

        Ok(())
    }

    async fn collect_metric(&mut self) -> anyhow::Result<()> {
        let rows = self
            .client
            .query(&self.system_metric_sql)
            .fetch_all::<(String, i64, String)>()
            .await?;

        for row in rows {
            if let Some(m) = self.metric_map.get_mut(&row.0) {
                m.set(row.1 as f64);
                continue;
            }

            let metric = prometheus::Gauge::new(&row.0, row.2)?;
            metric.set(row.1 as f64);

            DEFAULT_REGISTRY.register(Box::new(metric.clone())).ok();
            self.metric_map.insert(row.0, metric);
        }

        Ok(())
    }
}
